"""
数据加载模块
负责从CSV文件加载股票数据并进行预处理
"""

import pandas as pd
import numpy as np
from typing import Optional, List, Tuple
from .config import config, flags

class DataLoader:
    """数据加载器类"""
    
    def __init__(self):
        self.df: Optional[pd.DataFrame] = None
        self.raw_data: Optional[pd.DataFrame] = None
    
    def load_from_csv(self, file_path: str, date_format: str = '%Y%m%d') -> pd.DataFrame:
        """
        从CSV文件加载股票数据
        
        Args:
            file_path: CSV文件路径
            date_format: 日期格式，默认为'%Y%m%d'
            
        Returns:
            处理后的DataFrame
        """
        try:
            # 读取CSV文件
            self.raw_data = pd.read_csv(file_path)
            
            # 转换日期格式
            self.raw_data['trade_date'] = pd.to_datetime(
                self.raw_data['trade_date'], 
                format=date_format
            )
            
            # 按日期排序
            self.df = self.raw_data.sort_values('trade_date').reset_index(drop=True)
            
            print(f"成功加载数据，共 {len(self.df)} 条记录")
            print(self.df.head())
            
            return self.df
            
        except Exception as e:
            print(f"加载数据时出错: {e}")
            raise
    
    def get_data(self) -> pd.DataFrame:
        """获取处理后的数据"""
        if self.df is None:
            raise ValueError("数据尚未加载，请先调用 load_from_csv() 方法")
        return self.df
    
    def get_candlestick_data(self) -> List[Tuple[int, float, float, float, float]]:
        """
        转换为K线绘制所需的数据格式
        
        Returns:
            K线数据列表，格式为 [(索引, 开盘, 收盘, 最低, 最高), ...]
        """
        if self.df is None:
            raise ValueError("数据尚未加载")
        
        candlesticks = []
        for i, row in self.df.iterrows():
            candlesticks.append((
                i,  # x轴索引
                row['open'],   # 开盘价
                row['close'],  # 收盘价
                row['low'],    # 最低价
                row['high']    # 最高价
            ))
        
        return candlesticks
    
    def get_volume_data(self) -> Tuple[List[int], List[float], List[str]]:
        """
        获取成交量数据
        
        Returns:
            (x坐标列表, 成交量高度列表, 颜色列表)
        """
        if self.df is None:
            raise ValueError("数据尚未加载")
        
        volume_x = []
        volume_heights = []
        volume_colors = []
        
        for i, row in self.df.iterrows():
            volume_x.append(i)
            volume_heights.append(row['vol'])
            
            # 根据涨跌设置颜色
            if row['close'] >= row['open']:
                volume_colors.append(config.RISE_COLOR)  # 涨 - 红色
            else:
                volume_colors.append(config.FALL_COLOR)  # 跌 - 绿色
        
        return volume_x, volume_heights, volume_colors
    
    def get_price_range(self) -> Tuple[float, float]:
        """获取价格范围"""
        if self.df is None:
            raise ValueError("数据尚未加载")
        
        min_price = self.df['low'].min()
        max_price = self.df['high'].max()
        return min_price, max_price
    
    def get_volume_range(self) -> Tuple[float, float]:
        """获取成交量范围"""
        if self.df is None:
            raise ValueError("数据尚未加载")
        
        min_volume = self.df['vol'].min()
        max_volume = self.df['vol'].max()
        return min_volume, max_volume
    
    def get_date_ticks(self, max_ticks: int = 10) -> List[Tuple[int, str]]:
        """
        获取X轴日期刻度
        
        Args:
            max_ticks: 最大刻度数量
            
        Returns:
            日期刻度列表
        """
        if self.df is None:
            raise ValueError("数据尚未加载")
        
        xticks = [(i, str(d.date())) for i, d in enumerate(self.df['trade_date'])]
        step = max(1, len(xticks) // max_ticks)
        return xticks[::step]
    
    def get_data_by_index(self, index: int) -> Optional[dict]:
        """
        根据索引获取单条数据
        
        Args:
            index: 数据索引
            
        Returns:
            数据字典，如果索引无效返回None
        """
        if self.df is None or not (0 <= index < len(self.df)):
            return None
        
        row = self.df.iloc[index]
        return {
            'date': row['trade_date'],
            'open': row['open'],
            'high': row['high'],
            'low': row['low'],
            'close': row['close'],
            'vol': row['vol'],
            'amount': row['amount'] if 'amount' in row else 0,
            'change': row['close'] - row['open'],
            'pct_change': ((row['close'] - row['open']) / row['open'] * 100) if row['open'] != 0 else 0
        }
